/* 
* Copyright (C) 2010, John Asmuth

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
* 
*  $Revision$
*  $Date$
*  $Author$
*  $HeadURL$
* 
*/

package rlglue

import (
	"fmt"
	"math/rand"
	"strconv"
	"strings"
)

type Agent interface {
	AgentInit(taskString string)
	AgentStart(obs Observation) (act Action)
	AgentStep(reward float64, obs Observation) (act Action)
	AgentEnd(reward float64)
	AgentCleanup()
	AgentMessage(message string) (reply string)
}

type DefaultAgent struct{}

func (*DefaultAgent) AgentInit(string)                            {}
func (*DefaultAgent) AgentStart(Observation) (act Action)         { return }
func (*DefaultAgent) AgentStep(float64, Observation) (act Action) { return }
func (*DefaultAgent) AgentEnd(float64)                            {}
func (*DefaultAgent) AgentCleanup()                               {}
func (*DefaultAgent) AgentMessage(message string) (reply string) {
	if strings.HasPrefix(message, "seed ") {
		seedStr := message[len("seed "):]
		if seed, err := strconv.ParseInt(seedStr, 10, 64); err == nil {
			rand.Seed(seed)
		}
	}
	return
}

func LoadAgent(agent Agent) (err error) {
	ca := NewClientAgent(agent)
	if err = ca.AutoConnect(); err != nil {
		return
	}
	if err = ca.RunAgentEventLoop(); err != nil {
		return
	}
	err = ca.Close()
	return
}

type ClientAgent struct {
	agent Agent
	nw    *network
}

func NewClientAgent(agent Agent) *ClientAgent {
	return &ClientAgent{agent, nil}
}

func (ca *ClientAgent) onAgentInit() (err error) {
	taskSpec, err := ca.nw.GetString()
	ca.agent.AgentInit(taskSpec)
	ca.nw.ClearSendBuffer()
	ca.nw.PutInt(kAgentInit)
	ca.nw.PutInt(0)
	return
}

func (ca *ClientAgent) onAgentStart() (err error) {
	obs, err := ca.nw.GetObservation()
	act := ca.agent.AgentStart(obs)
	size := ca.nw.SizeOfAction(act)
	ca.nw.ClearSendBuffer()
	ca.nw.PutInt(kAgentStart)
	ca.nw.PutInt(size)
	ca.nw.PutAction(act)
	return
}

func (ca *ClientAgent) onAgentStep() (err error) {
	r, err := ca.nw.GetDouble()
	if err != nil {
		return
	}
	obs, err := ca.nw.GetObservation()
	if err != nil {
		return
	}
	act := ca.agent.AgentStep(r, obs)
	size := ca.nw.SizeOfAction(act)
	ca.nw.ClearSendBuffer()
	ca.nw.PutInt(kAgentStep)
	ca.nw.PutInt(size)
	ca.nw.PutAction(act)
	return
}

func (ca *ClientAgent) onAgentEnd() (err error) {
	r, err := ca.nw.GetDouble()
	if err != nil {
		return
	}
	ca.agent.AgentEnd(r)
	ca.nw.ClearSendBuffer()
	ca.nw.PutInt(kAgentEnd)
	ca.nw.PutInt(0)
	return
}

func (ca *ClientAgent) onAgentCleanup() (err error) {
	ca.agent.AgentCleanup()
	ca.nw.ClearSendBuffer()
	ca.nw.PutInt(kAgentCleanup)
	ca.nw.PutInt(0)
	return
}

func (ca *ClientAgent) onAgentMessage() (err error) {
	message, err := ca.nw.GetString()
	if err != nil {
		return
	}
	reply := ca.agent.AgentMessage(message)
	ca.nw.ClearSendBuffer()
	ca.nw.PutInt(kAgentMessage)
	ca.nw.PutInt(int32(len(reply)) + kIntSize)
	ca.nw.PutString(reply)
	return
}

func (ca *ClientAgent) Connect(host string, port int) (err error) {
	ca.nw = new(network)
	err = ca.nw.Connect(host, port)
	if err != nil {
		return
	}
	ca.nw.ClearSendBuffer()
	ca.nw.PutInt(kAgentConnection)
	ca.nw.PutInt(0)
	err = ca.nw.Send()
	return
}

func (ca *ClientAgent) AutoConnect() (err error) {
	host, port := GetEnvHostPort()
	err = ca.Connect(host, port)
	return
}

func (ca *ClientAgent) Close() error {
	return ca.nw.Close()
}

func (ca *ClientAgent) RunAgentEventLoop() (err error) {
	var agentState int32 = 0
	var recvSize int32 = 0

	for agentState != kRLTerm {
		ca.nw.ClearRecvBuffer()

		// We may have received the header and part of the payload
		// We need to keep track of how much of the payload 
		recvSize, err = ca.nw.Recv(2 * kIntSize)

		if err != nil {
			return
		}

		recvSize -= 2 * kIntSize

		agentState, err = ca.nw.GetInt()
		if err != nil {
			return
		}
		var dataSize int32
		dataSize, err = ca.nw.GetInt()
		if err != nil {
			return
		}

		//fmt.Printf("agent state: %d\n", agentState)

		remaining := dataSize - recvSize
		if remaining < 0 {
			print("Remaining was less than 0!\n")
			remaining = 0
		}

		_, err = ca.nw.Recv(remaining)
		if err != nil {
			return
		}

		// Already read the header, discard it
		//ca.nw.GetInt()
		//ca.nw.GetInt()

		switch agentState {
		case kAgentInit:
			ca.onAgentInit()
		case kAgentStart:
			ca.onAgentStart()
		case kAgentStep:
			ca.onAgentStep()
		case kAgentEnd:
			ca.onAgentEnd()
		case kAgentCleanup:
			ca.onAgentCleanup()
		case kAgentMessage:
			ca.onAgentMessage()
		case kRLTerm:
			break
		default:
			panic(fmt.Sprintf("Unknown message type: %d", agentState))
		}

		err = ca.nw.Send()
		if err != nil {
			return
		}
	}

	return
}
